package com.yanyeori.framework.datasource.switcher;

import com.alibaba.ttl.threadpool.TtlExecutors;
import com.yanyeori.framework.core.thread.ThreadFactoryManager;
import com.yanyeori.framework.datasource.config.DatabaseContextHolder;
import com.yanyeori.framework.datasource.configprop.DataSourceConfigProps;
import lombok.extern.slf4j.Slf4j;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.util.StopWatch;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.*;

/**
 * 数据源存活检测
 *
 * @author chenkuan 2021/8/19
 */
@Slf4j
public class SurviveEviction {

    private final JdbcTemplate jdbcTemplate;
    private final DataSourceConfigProps config;

    private static final String QUERY_SQL = "SELECT 1 FROM DUAL";

    private final Map<String, SurviveStatus> SURVIVE_STATUS_MAP = new HashMap<>();

    private final ExecutorService executorService = TtlExecutors.getTtlExecutorService(getExecutorService());

    public SurviveEviction(JdbcTemplate jdbcTemplate, DataSourceConfigProps config) {
        this.jdbcTemplate = jdbcTemplate;
        this.config = config;
    }

    /**
     * 数据源存活检测
     */
    public SurviveStatus survive(String dbKey) {
        //检查距离上次检查时间间隔
        long nowMillis = System.currentTimeMillis();
        boolean containsKey = SURVIVE_STATUS_MAP.containsKey(dbKey);
        SurviveStatus surviveStatus = SURVIVE_STATUS_MAP.computeIfAbsent(dbKey, k -> new SurviveStatus(true, nowMillis, 0));
        if (containsKey && nowMillis - surviveStatus.getLastSurviveTime() < config.getConfig().getDbEvictionTimeMillis()) {
            //该数据源两次检测时间小于设定时间，本次不检测
            return surviveStatus;
        }

        Future<Boolean> future = executorService.submit(() -> {
            StopWatch stopWatch = new StopWatch();
            stopWatch.start();
            try {
                DatabaseContextHolder.setDBKey(dbKey);
                jdbcTemplate.execute(QUERY_SQL);
                stopWatch.stop();
                log.debug("Datasource [{}] is successfully detected and takes {} ms", dbKey, stopWatch.getLastTaskTimeMillis());
                averageTime(surviveStatus, stopWatch.getLastTaskTimeMillis());
                return true;
            } catch (Exception e) {
                stopWatch.stop();
                log.warn("Datasource [{}] detection failed and takes {} ms", dbKey, stopWatch.getLastTaskTimeMillis());
                return false;
            }
        });

        try {
            long dbEvictionTimeoutMillis = config.getConfig().getDbEvictionTimeoutMillis();
            long start = System.currentTimeMillis();
            while (!future.isDone()) {
                TimeUnit.MILLISECONDS.sleep(surviveStatus.getElapsedTime());
                if (System.currentTimeMillis() - start > dbEvictionTimeoutMillis) {
                    break;
                }
            }
            boolean done = future.isDone();
            if (done && future.get()) {
                return surviveSuccess(surviveStatus);
            } else {
                if (!done) {
                    future.cancel(true);
                }
                return surviveFailed(surviveStatus);
            }
        } catch (Exception e) {
            return surviveFailed(surviveStatus);
        }
    }

    private SurviveStatus surviveSuccess(SurviveStatus surviveStatus) {
        surviveStatus.setSurvive(true);
        surviveStatus.setLastSurviveTime(System.currentTimeMillis());
        surviveStatus.setErrorTimes(0);
        return surviveStatus;
    }

    private SurviveStatus surviveFailed(SurviveStatus surviveStatus) {
        surviveStatus.setSurvive(false);
        surviveStatus.setLastSurviveTime(System.currentTimeMillis());
        surviveStatus.setErrorTimes(surviveStatus.getErrorTimes() + 1);
        return surviveStatus;
    }

    /**
     * 计算检测耗时均值
     */
    private void averageTime(SurviveStatus surviveStatus, long lastTaskTimeMillis) {
        surviveStatus.setElapsedTime(Math.round((double) (surviveStatus.getElapsedTime() * 4 + lastTaskTimeMillis) / 5));
    }

    private ExecutorService getExecutorService() {
        return new ThreadPoolExecutor(
                1,
                8,
                30L,
                TimeUnit.SECONDS,
                new SynchronousQueue<>(),
                ThreadFactoryManager.namedThreadFactory("db-check"),
                new ThreadPoolExecutor.DiscardPolicy());
    }
}
